
import numpy as np
from math import sqrt, floor, ceil
import matplotlib.pyplot as plt

def wma(series, period):
    """
    加权移动平均 (Weighted Moving Average)

    参数:
    - series: 价格序列 (numpy数组或pandas Series)
    - period: 加权平均周期

    返回:
    - WMA值序列
    """
    if len(series) < period:
        return np.full(len(series), np.nan)

    weights = np.arange(1, period + 1)  # 权重: 1, 2, 3, ..., period
    weights_sum = weights.sum()

    result = np.full(len(series), np.nan)

    for i in range(period - 1, len(series)):
        # 取最近period个值
        window = series[i - period + 1: i + 1]
        # 加权平均: sum(价格 * 权重) / 权重总和
        wma_value = np.sum(window * weights) / weights_sum
        result[i] = wma_value

    return result

def tdx_double_wma(close_prices, n):
    """
    实现通达信双重加权移动平均公式

    参数:
    - close_prices: 收盘价序列
    - n: 基础周期参数

    返回:
    - 双重WMA指标值
    """
    # 转换为numpy数组
    close = np.asarray(close_prices, dtype=np.float64)

    # 计算中间周期
    period1 = round(n / 2)  # ROUND(N/2)
    period2 = round(sqrt(n))  # ROUND(SQRT(N))

    # 第一步: 计算两个WMA
    wma_short = wma(close, period1)  # WMA(C, ROUND(N/2))
    wma_long = wma(close, n)  # WMA(C, N)

    # 第二步: 计算差值 2*短期 - 长期
    diff = 2 * wma_short - wma_long

    # 第三步: 对差值进行最终WMA平滑
    result = wma(diff, period2)  # WMA(DIFF, ROUND(SQRT(N)))

    return result





def visualize_parameter_impact(close_prices, n_values=[21,88]):
    # 创建测试数据并可视化
    np.random.seed(42)
    test_prices = 100 + np.cumsum(np.random.randn(100) * 2)
    close_prices =test_prices
    """
    可视化不同N值对指标的影响
    """
    plt.figure(figsize=(12, 8))

    # 绘制价格
    plt.plot(close_prices, label='收盘价', color='black', linewidth=2, alpha=0.7)

    # 计算并绘制不同N值的指标
    colors = ['red', 'blue', 'green', 'orange']

    for i, n in enumerate(n_values):
        # 计算指标
        indicator = tdx_double_wma_single_param(close_prices, n)

        # 只绘制有效部分
        valid_idx = ~np.isnan(indicator)
        if np.any(valid_idx):
            # 显示参数分解信息
            p1 = int(round(n / 2))
            p3 = int(round(sqrt(n)))

            plt.plot(np.where(valid_idx)[0], indicator[valid_idx],
                     label=f'N={n} (周期:{p1},{n},{p3})',
                     color=colors[i % len(colors)], linewidth=1.5)

    plt.title('不同N值对双重WMA指标的影响')
    plt.xlabel('时间周期')
    plt.ylabel('指标值')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.show()